#!/usr/bin/env python
#-*- coding: utf-8 -*-

import sys
import commands
import json


def _get_gpu_num():
	cmd = 'nvidia-smi -L |wc -l'
	cmd_status, gpu_num = commands.getstatusoutput(cmd)
	if cmd_status == 0:
		return gpu_num
	else:
		return 0

def gpu_discovery():
	gpu_num = int(_get_gpu_num())
	gpu_list = []
	gpu_json_list = []
	for i in range(gpu_num):
		gpu_list.append('GPU %r' % i)
		gpu_json_list.append({"{#NVDGPU}": "GPU%r" % i})
	gpu_json_fmt = json.dumps({'data': gpu_json_list}, sort_keys=True, indent=4)
		
	return gpu_json_fmt
	

def _get_gpu_info(gpu_order):
	cmd_get_memory_used = 'nvidia-smi -i %r -q -d MEMORY|grep -v N/A |grep  Used|awk \'{print $3}\'' % gpu_order
	cmd_get_memory_total = 'nvidia-smi -i %r -q -d MEMORY|grep -v N/A |grep  Total|awk \'{print $3}\'' % gpu_order
	cmd_get_memory_free = 'nvidia-smi -i %r -q -d MEMORY|grep -v N/A |grep  Free|awk \'{print $3}\'' % gpu_order
	cmd_get_temperature = 'nvidia-smi -i 0 -q -d TEMPERATURE|grep Current|awk \'{print $5}\''
	memory_used = int(commands.getoutput(cmd_get_memory_used))
	memory_total = int(commands.getoutput(cmd_get_memory_total))
	memory_free = int(commands.getoutput(cmd_get_memory_free))
	temperature = int(commands.getoutput(cmd_get_temperature))
	memory_use_pct = round(memory_used / float(memory_total) * 100, 2)
	return [memory_total, memory_used, memory_use_pct, temperature]

if len(sys.argv) == 2:
	if sys.argv[1] == 'discovery':
		gpu_num = gpu_discovery()
		print gpu_num
elif len(sys.argv) == 3:
	gpu_info = _get_gpu_info(sys.argv[1].split('GPU')[1])
	if sys.argv[2] == 'total':
		print gpu_info[0]
	elif sys.argv[2] == 'usage':
		print gpu_info[1]
	elif sys.argv[2] == 'pct':
		print gpu_info[2]
	elif sys.argv[2] == 'temp':
		print gpu_info[3]

	
